import os
import random

import torch
from torch.utils.tensorboard import SummaryWriter

from config.device_config import get_device
from model.loss.loss_point_picker import torch_d_P_Ps
from model.mesh_dataset import MeshDataset
from model.model_point_picker import ModelPointPicker

writer = SummaryWriter()
torch.manual_seed(42)

save_model_every = 1000       # batch
max_epoch = 200*save_model_every
min_number_points = 1
max_number_points = 25_000
PROJECT_ROOT = r"C:\Users\47306\PycharmProjects\Mesh-simplification-GNN"

from torchinfo import summary
device  = get_device()
def train():
    torch_dataset = MeshDataset(os.path.join(PROJECT_ROOT, "3d_models/stl/"))
    print(f'device...:{device}')
    model_point_picker = ModelPointPicker().to(device)
    print(model_point_picker)
    optimizer = torch.optim.Adam(model_point_picker.parameters(), lr=1e-3)
    
    for epoch in range(max_epoch+1):
        d_P_Ps_epoch = 0.0
        simplification_rate = random.random()
        
        for i, (torch_graph, _) in enumerate(torch_dataset):
            # 将 torch_graph 移到 device 上
            torch_graph = torch_graph.to(device)
            target_number_point = int(max(min_number_points, min(len(torch_graph.x)*simplification_rate, max_number_points)))
            score_original_points, generated_graph_nodes = model_point_picker(target_number_point, torch_graph)
            print(f'score_original_points:{score_original_points}')

            def process_input(torch_graph):
                node_features = torch_graph.x.to(device)
                edge_index = torch_graph.edge_index.to(device)
                return {"x": node_features, "edge_index": edge_index}

            processed_input = process_input(torch_graph)
            summary(model_point_picker, input_data=(target_number_point, processed_input), device=device)

            d_P_Ps = torch_d_P_Ps(score_original_points, torch_graph.x, generated_graph_nodes, simplification_rate)
            d_P_Ps_epoch += d_P_Ps
        
        d_P_Ps_epoch.backward()
        optimizer.step()
        optimizer.zero_grad()

        writer.add_scalar('Loss/d_P_Ps', d_P_Ps_epoch, epoch)
        print('EPOCH : ', epoch, '   |   SIMPLIFICATION : ', str(int(simplification_rate*100)),'%   |   LOSS : ', d_P_Ps.data.item())

        if epoch%save_model_every==0:
            path = os.path.join(PROJECT_ROOT,'save_models/point_picker/10')
            if not os.path.exists(path):
                os.makedirs(path)
            torch.save(model_point_picker.state_dict(),os.path.join(PROJECT_ROOT,path+'/'+str(epoch)+'.pt'))

train()